# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# PyTorch (https://pytorch.org/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Base cell for bais potential"""

from inspect import signature

import torch
from torch import Tensor
from torch import nn
import torch.nn.functional as F

from .bias import Bias
from ...function import get_arguments, get_tensor
from ...function.units import Length, Energy


class SphericalRestrict(Bias):
    r"""Basic cell for bias potential

    Math:

    .. Math::

        V(R) = k \log{\left ( 1 + e^{\frac{|R - R_0| - r_0}{\sigma}} \right )}

    Args:

        radius (float):         Radius of sphere (r_0).

        center (Tensor):        Coordinate of the center of sphere (R_0).

        force_constant (float): Force constant of the bias potential(k). Default: Energy(500, 'kj/mol')

        depth (float):          Wall depth of the restriction (\sigma). Default: Length(0.01, 'nm')

        length_unit (str):      Length unit. If None is given, it will be assigned with the global length unit.
                                Default: None

        energy_unit (str):      Energy unit. If None is given, it will be assigned with the global energy unit.
                                Default: None

        use_pbc (bool):         Whether to use periodic boundary condition.

        name (str):             Name of the bias potential. Default: 'spherical'

    Supported Platforms:

        ``CPU`` ``GPU``

    """
    def __init__(self,
                 radius: float,
                 center: Tensor = 0,
                 force_constant: float = Energy(500, 'kj/mol'),
                 depth: float = Length(0.01, 'nm'),
                 length_unit: str = None,
                 energy_unit: str = None,
                 use_pbc: bool = None,
                 name: str = 'spherical',
                 **kwargs,
                 ):

        super().__init__(
            name=name,
            length_unit=length_unit,
            energy_unit=energy_unit,
            use_pbc=use_pbc,
        )
        self._kwargs = get_arguments(locals(), kwargs)

        self.radius = get_tensor(radius, dtype=torch.float32)
        self.center = get_tensor(center, dtype=torch.float32)

        if isinstance(force_constant, Energy):
            force_constant = force_constant(self.units)
        self.force_constant = get_tensor(force_constant, dtype=torch.float32)

        if isinstance(depth, Length):
            depth = depth(self.units)
        self.depth = get_tensor(depth, dtype=torch.float32)

    def calc_colvar(self,
                    coordinate: Tensor,
                    neighbour_index: Tensor = None,
                    neighbour_mask: Tensor = None,
                    neighbour_vector: Tensor = None,
                    neighbour_distance: Tensor = None,
                    pbc_box: Tensor = None
                    ):
        """calculate the value of collective variables"""
        # (B, A, D) - (D)
        vector = coordinate - self.center
        # (B, A) <- (B, A, D)
        distance = torch.norm(vector, dim=-1)
        return distance

    def calc_bias(self, colvar: Tensor) -> Tensor:
        """calculate the value of bias potential as the function of collective variables"""
        diff = colvar - self.radius
        bias = self.force_constant * torch.log1p(torch.exp(diff/self.depth))
        # (B, 1) <- (B, A)
        return torch.sum(bias, -1, keepdim=True)

    def print_info(self) -> dict:
        """Return basic information about the spherical restriction bias potential
        
        Returns:
            dict: Dictionary containing spherical restriction information
        """
        info = super().print_info()  # 获取基类信息
        info.update({
            "Radius": f"{self.radius.item():.3f} {self.units.length_unit}",
            "Center": f"{self.center.tolist()}",
            "Force constant": f"{self.force_constant.item():.3f} {self.units.energy_unit}",
            "Wall depth": f"{self.depth.item():.3f} {self.units.length_unit}"
        })
        return info
